# calibration.py
#!/usr/bin/env python
import cv2
import yaml
import glob
import numpy as np
from camera import *

# Defining the dimensions of checkerboard
CHECKERBOARD = (5, 8)
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
SQUARE_SIZE = 50 # 打印的格子边长，单位 mm

# Defining the world coordinates for 3D points
objp = np.zeros((CHECKERBOARD[0] * CHECKERBOARD[1], 3), np.float32)
objp[:,:2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2)
objp = objp*SQUARE_SIZE

class Calibration():
    """Camera calibration parameters
    """
    def __init__(self):
        self.mtx = np.array(1)
        self.dist = np.array(1)
        self.camera_u = 0
        self.camera_v = 0

    def toString(self):
        print("Camera matrix : \n")
        print(self.mtx)
        print("dist : \n")
        print(self.dist)
        # print("rvecs : \n")
        # print(self.rvecs)
        # print("tvecs : \n")
        # print(self.tvecs)

    def load(self, fn:str):
        # Load calibration parameters from file
        with open(fn, "r") as file:
            parameter = yaml.load(file.read(), Loader=yaml.Loader)
            mtx = parameter['camera_matrix']
            dist = parameter['dist_coeff']
            self.camera_u = parameter['camera_u']
            self.camera_v = parameter['camera_v']
            self.mtx = np.array(mtx)
            self.dist = np.array(dist)

    def save(self, fn:str):
        # Save calibration parameters to file
        #存取标定数据
        mtx_yaml = self.mtx.tolist()
        dist_yaml = self.dist.tolist()
        data = {"information":"Camera calibration parameters",
                "camera_matrix":mtx_yaml,
                "dist_coeff":dist_yaml,
                "camera_u":self.camera_u,
                "camera_v":self.camera_v }
        with open(fn,"w") as file:
            yaml.dump(data,file)

    def run(self, img_dir:str):
        # Creating vector to store vectors of 3D points for each checkerboard image
        obj_points = []
        # Creating vector to store vectors of 2D points for each checkerboard image
        img_points = []

        # 对指定文件夹的图片列表进行批处理校准计算
        # 加载pic文件夹下所有的jpg图像
        images = glob.glob(f'{img_dir}/*.jpg')
        imgSize = (0,0)
        for fname in images:
            img = cv2.imread(fname)
            # 获取画面中心点
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            # print(gray.shape[::-1])
            imgSize = gray.shape[::-1]
            # Find the chess board corners
            # If desired number of corners are found in the image then ret = true
            ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD, None)
            """
            If desired number of corner are detected,
            we refine the pixel coordinates and display
            them on the images of checker board
            """
            if ret == True:
                obj_points.append(objp)
                # refining pixel coordinates for given 2d points.
                corners2 = cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria)
                if corners2.any():
                    img_points.append(corners2)
                else:
                    img_points.append(corners)

        """
        Performing camera calibration by passing the value of known 3D points (obj_points)
        and corresponding pixel coordinates of the detected corners (img_points)
        """
        ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(obj_points, img_points, imgSize, None, None)
        if ret:
            self.mtx = mtx
            self.dist = dist
            self.rvecs = rvecs
            self.tvecs = tvecs
            self.toString()

if __name__ == '__main__':
    calib = Calibration()
    calib.run('D:/MyProject/Vislander')
    calib.save('./标定文件.yaml')
